#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
污染类型数据库表创建脚本
手动创建数据库表结构，避免使用alembic
"""

import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from core.database import engine
from models.pollution_type import PollutionType, PollutionTypeRelation, PollutionTypeExample
from core.logging_config import get_logger

logger = get_logger("pollution_type_migration")

def create_pollution_type_tables():
    """创建污染类型相关的数据库表"""
    try:
        logger.info("开始创建污染类型相关表...")
        
        # 创建表
        PollutionType.__table__.create(engine, checkfirst=True)
        PollutionTypeRelation.__table__.create(engine, checkfirst=True)
        PollutionTypeExample.__table__.create(engine, checkfirst=True)
        
        logger.info("污染类型相关表创建完成")
        return True
        
    except Exception as e:
        logger.error(f"创建污染类型表失败: {str(e)}")
        return False

def drop_pollution_type_tables():
    """删除污染类型相关的数据库表（谨慎使用）"""
    try:
        logger.info("开始删除污染类型相关表...")
        
        # 按照依赖关系反向删除
        PollutionTypeExample.__table__.drop(engine, checkfirst=True)
        PollutionTypeRelation.__table__.drop(engine, checkfirst=True)
        PollutionType.__table__.drop(engine, checkfirst=True)
        
        logger.info("污染类型相关表删除完成")
        return True
        
    except Exception as e:
        logger.error(f"删除污染类型表失败: {str(e)}")
        return False

def check_tables_exist():
    """检查表是否存在"""
    try:
        from sqlalchemy import text
        with engine.connect() as conn:
            # 检查pollution_types表
            result = conn.execute(text("SELECT name FROM sqlite_master WHERE type='table' AND name='pollution_types'"))
            pollution_types_exists = result.fetchone() is not None
            
            # 检查pollution_type_relations表
            result = conn.execute(text("SELECT name FROM sqlite_master WHERE type='table' AND name='pollution_type_relations'"))
            relations_exists = result.fetchone() is not None
            
            # 检查pollution_type_examples表
            result = conn.execute(text("SELECT name FROM sqlite_master WHERE type='table' AND name='pollution_type_examples'"))
            examples_exists = result.fetchone() is not None
            
            logger.info(f"表存在状态: pollution_types={pollution_types_exists}, relations={relations_exists}, examples={examples_exists}")
            
            return {
                "pollution_types": pollution_types_exists,
                "pollution_type_relations": relations_exists,
                "pollution_type_examples": examples_exists
            }
            
    except Exception as e:
        logger.error(f"检查表存在状态失败: {str(e)}")
        return {}

if __name__ == "__main__":
    print("污染类型数据库表管理脚本")
    print("=" * 50)
    
    # 检查当前表状态
    tables_status = check_tables_exist()
    print(f"当前表状态: {tables_status}")
    
    # 询问用户操作
    print("\n请选择操作:")
    print("1. 创建表（如果不存在）")
    print("2. 重新创建表（会删除现有数据）")
    print("3. 仅检查表状态")
    print("4. 退出")
    
    choice = input("\n请输入选项 (1-4): ").strip()
    
    if choice == "1":
        success = create_pollution_type_tables()
        if success:
            print("表创建成功！")
        else:
            print("表创建失败，请查看日志。")
            
    elif choice == "2":
        confirm = input("警告：这将删除所有现有数据！确定要继续吗？(y/N): ").strip().lower()
        if confirm == 'y':
            drop_pollution_type_tables()
            success = create_pollution_type_tables()
            if success:
                print("表重新创建成功！")
            else:
                print("表重新创建失败，请查看日志。")
        else:
            print("操作已取消。")
            
    elif choice == "3":
        print("表状态检查完成。")
        
    elif choice == "4":
        print("退出脚本。")
        
    else:
        print("无效选项。")